import sys
import uuid

from clu import metric_writers
import jax
from brax import envs
from brax.io import html
from brax.io import model
from brax.training import apg
from brax.training import es
from brax.training import ppo
from brax.training import sac
from robotisgp import Robotisgp

logdir = sys.argv[1]
env_fn = Robotisgp
env = env_fn()
state = env.reset(jax.random.PRNGKey(0))
# output an episode trajectory
qps = []
jit_step_fn = jax.jit(env.step)
rng = jax.random.PRNGKey(0)
T = 1000
t = 0
while not state.done and (t<T):
  qps.append(state.qp)
  key, rng = jax.random.split(rng)
  act = jax.random.uniform(key,(env.action_size, ))
  state = jit_step_fn(state, act)
  t += 1

html_path = f'{logdir}/trajectory_{uuid.uuid4()}.html'
html.save_html(html_path, env.sys, qps)
